#include <iostream>
#include <cstring>
#include <thread>
#include <vector>
#include <atomic>  // 添加原子类型支持
#include <chrono>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>

const char* SERVER_ADDRESS = "127.0.0.1"; // 替换为你的服务器地址
const int SERVER_PORT = 8000;              // 替换为你的服务器端口
const int MESSAGE_SIZE = 1000;             // 每个消息发送 1000 字节
const int NUM_MESSAGES = 1;               // 每个连接发送 几 次消息
const int SLEEP_DURATION_MS = 50;          // 每次发送消息后等待 50 毫秒
const int num_clients = 1;               // 设置客户端数量

// 统计成功接收的响应数量
std::atomic<int> success_count(0); // 原子计数器，用于线程安全地统计成功的连接

void tcp_client(int thread_id) {
    int sock;
    struct sockaddr_in server_addr;
    int total_sent = 0;         // 统计发送的总字节数
    int total_received = 0;     // 统计接收的总字节数

    // 创建 socket
    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        std::cerr << "Thread " << thread_id << ": Socket creation error" << std::endl;
        return;
    }

    // 配置服务器地址结构
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(SERVER_PORT);

    // 转换地址
    if (inet_pton(AF_INET, SERVER_ADDRESS, &server_addr.sin_addr) <= 0) {
        std::cerr << "Thread " << thread_id << ": Invalid address/ Address not supported" << std::endl;
        close(sock);
        return;
    }

    // 连接到服务器
    if (connect(sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        std::cerr << "Thread " << thread_id << ": Connection to server failed" << std::endl;
        close(sock);
        return;
    }

    // 发送消息并接收响应，共发送 10 次
    for (int i = 0; i < NUM_MESSAGES; ++i) {
        char message[MESSAGE_SIZE];
        std::memset(message, 'A', MESSAGE_SIZE); // 填充消息

        // 发送消息
        int sent_bytes = send(sock, message, MESSAGE_SIZE, 0);
        if (sent_bytes < 0) {
            std::cerr << "Thread " << thread_id << ": Failed to send message" << std::endl;
            break;
        }
        total_sent += sent_bytes;  // 累计发送的字节数

        // 接收响应
        char buffer[1024] = {0};
        int bytes_received = recv(sock, buffer, sizeof(buffer), 0);
        if (bytes_received < 0) {
            std::cerr << "Thread " << thread_id << ": Failed to receive message" << std::endl;
            break;
        }
        total_received += bytes_received;  // 累计接收的字节数

        // 每次发送后等待 50 毫秒
        std::this_thread::sleep_for(std::chrono::milliseconds(SLEEP_DURATION_MS));
    }

    // 判断是否成功（发送和接收的数据量是否一致）
    if (total_sent == total_received) {
        std::cout << "Thread " << thread_id << ": Success! Sent: " << total_sent << " bytes, Received: " << total_received << " bytes" << std::endl;
        success_count++;  // 统计成功的连接数
    } else {
        std::cout << "Thread " << thread_id << ": Mismatch! Sent: " << total_sent << " bytes, Received: " << total_received << " bytes" << std::endl;
    }

    // 关闭 socket
    close(sock);
}

int main() {
    
    std::vector<std::thread> threads;

    // 启动多个线程来测试服务器
    for (int i = 0; i < num_clients; ++i) {
        threads.emplace_back(tcp_client, i);
    }

    // 等待所有线程完成
    for (auto& thread : threads) {
        thread.join();
    }

    // 输出成功率
    std::cout << "Success Rate: " << success_count.load() << "/" << num_clients << std::endl;

    return 0;
}
